// Copyright (C) Kumo inc. and its affiliates.
// Author: Jeff.li lijippy@163.com
// All rights reserved.
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU Affero General Public License as published
// by the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU Affero General Public License for more details.
//
// You should have received a copy of the GNU Affero General Public License
// along with this program.  If not, see <https://www.gnu.org/licenses/>.
//
#include <pollux/expression/switch_expr.h>
#include <pollux/expression/boolean_mix.h>
#include <pollux/expression/constant_expr.h>
#include <pollux/expression/field_reference.h>
#include <pollux/expression/scoped_var_setter.h>

namespace kumo::pollux::exec {
    namespace {
        bool hasElseClause(const std::vector<ExprPtr> &inputs) {
            return inputs.size() % 2 == 1;
        }
    } // namespace

    SwitchExpr::SwitchExpr(
        TypePtr type,
        const std::vector<ExprPtr> &inputs,
        bool inputsSupportFlatNoNullsFastPath)
        : SpecialForm(
              std::move(type),
              inputs,
              "switch",
              hasElseClause(inputs) && inputsSupportFlatNoNullsFastPath,
              false /* trackCpuUsage */),
          numCases_{inputs_.size() / 2},
          hasElseClause_{hasElseClause(inputs_)} {
        std::vector<TypePtr> inputTypes;
        inputTypes.reserve(inputs_.size());
        std::transform(
            inputs_.begin(),
            inputs_.end(),
            std::back_inserter(inputTypes),
            [](const ExprPtr &expr) { return expr->type(); });

        // Apply type checking.
        auto typeExpected = resolveType(inputTypes);
        POLLUX_CHECK(
            *typeExpected == *this->type(),
            "Switch expression type different than then clause. Expected {} but got Actual {}.",
            typeExpected->toString(),
            this->type()->toString());
    }

    void SwitchExpr::evalSpecialForm(
        const SelectivityVector &rows,
        EvalCtx &context,
        VectorPtr &finalResult) {
        VectorPtr localResult;
        LocalSelectivityVector remainingRows(context, rows);

        LocalSelectivityVector thenRows(context);

        // SWITCH: fix finalSelection at "rows" unless already fixed
        ScopedFinalSelectionSetter scopedFinalSelectionSetter(context, &rows);
        if (propagatesNulls_) {
            // If propagates nulls, we load lazies before conditions so that we can
            // avoid errors for null rows. Null propagation is on only if all thens and
            // else load access the same vectors, so there is no extra loading.
            auto &remaining = *remainingRows.get();
            for (auto *field: distinctFields_) {
                context.ensureFieldLoaded(field->index(context), remaining);
                const auto &vector = context.getField(field->index(context));
                if (vector->may_have_nulls()) {
                    LocalDecodedVector decoded(context, *vector, remaining);
                    add_nulls(remaining, decoded->nulls(&remaining), context, localResult);
                    remaining.deselectNulls(
                        decoded->nulls(&remaining), remaining.begin(), remaining.end());
                }
            }
        }

        VectorPtr condition;
        const uint64_t *values;

        for (auto i = 0; i < numCases_; i++) {
            context.releaseVector(condition);

            if (!remainingRows.get()->hasSelections()) {
                break;
            }

            // evaluate the case condition
            inputs_[2 * i]->eval(*remainingRows.get(), context, condition);

            if (context.errors()) {
                context.deselectErrors(*remainingRows);
                if (!remainingRows->hasSelections()) {
                    break;
                }
            }

            const auto booleanMix = getFlatBool(
                condition.get(),
                *remainingRows.get(),
                context,
                &tempValues_,
                nullptr,
                true,
                &values,
                nullptr);
            switch (booleanMix) {
                case BooleanMix::kAllTrue:
                    inputs_[2 * i + 1]->eval(*remainingRows.get(), context, localResult);
                    remainingRows->clearAll();
                    continue;
                case BooleanMix::kAllNull:
                case BooleanMix::kAllFalse:
                    continue;
                default: {
                    thenRows.get(remainingRows->end(), false);
                    bits::andBits(
                        thenRows.get()->asMutableRange().bits(),
                        remainingRows.get()->asRange().bits(),
                        values,
                        0,
                        remainingRows->end());
                    thenRows.get()->updateBounds();

                    if (thenRows.get()->hasSelections()) {
                        inputs_[2 * i + 1]->eval(*thenRows.get(), context, localResult);
                        remainingRows.get()->deselect(*thenRows.get());
                    }
                }
            }
        }

        // Evaluate the "else" clause.
        if (remainingRows.get()->hasSelections()) {
            if (hasElseClause_) {
                inputs_.back()->eval(*remainingRows.get(), context, localResult);
            } else {
                context.ensure_writable(*remainingRows.get(), type(), localResult);

                // fill in nulls for remainingRows
                remainingRows.get()->applyToSelected(
                    [&](auto row) { localResult->set_null(row, true); });
            }
        }

        // Some rows may have not been evaluated by any then or else clause because
        // a condition threw an error on these rows. We set those to nulls to make
        // sure the result vector is addressable at those indices.
        if (context.errors()) {
            // TODO: Fix decoding function vector issue #6269.
            if (type()->kind() != TypeKind::FUNCTION) {
                LocalSelectivityVector nonErrorRows(context, rows);
                context.deselectErrors(*nonErrorRows);
                add_nulls(rows, nonErrorRows->asRange().bits(), context, localResult);
            }
        }
        // TODO: Fix evaluate lambda expression return vector of size 0 issue #6270.
        if (type()->kind() != TypeKind::FUNCTION) {
            POLLUX_CHECK(localResult && localResult->size() >= rows.end());
        }

        context.moveOrCopyResult(localResult, rows, finalResult);
    }

    // This is safe to call only after all metadata is computed for input
    // expressions.
    void SwitchExpr::computePropagatesNulls() {
        // The "switch" expression propagates nulls when all of the following
        // conditions are met:
        // - All "then" clauses and optional "else" clause propagate nulls.
        // - All "then" clauses and optional "else" clause use the same inputs.
        // - All "condition" clauses use a subset of "then"/"else" inputs.

        for (auto i = 0; i < numCases_; i += 2) {
            if (!inputs_[i + 1]->propagatesNulls()) {
                propagatesNulls_ = false;
                return;
            }
        }

        const auto &firstThenFields = inputs_[1]->distinctFields();
        for (auto i = 0; i < numCases_; ++i) {
            const auto &condition = inputs_[i * 2];
            const auto &thenClause = inputs_[i * 2 + 1];
            if (!Expr::isSameFields(firstThenFields, thenClause->distinctFields())) {
                propagatesNulls_ = false;
                return;
            }

            if (!Expr::isSubsetOfFields(condition->distinctFields(), firstThenFields)) {
                propagatesNulls_ = false;
                return;
            }
        }

        if (hasElseClause_) {
            const auto &elseClause = inputs_.back();
            if (!elseClause->propagatesNulls()) {
                propagatesNulls_ = false;
                return;
            }
            if (!Expr::isSameFields(firstThenFields, elseClause->distinctFields())) {
                propagatesNulls_ = false;
                return;
            }
        }

        propagatesNulls_ = true;
    }

    // static
    TypePtr SwitchExpr::resolveType(const std::vector<TypePtr> &argTypes) {
        POLLUX_CHECK_GT(
            argTypes.size(),
            1,
            "Switch statements expect at least 2 arguments, received {}",
            argTypes.size());
        // Type structure is [cond1Type, then1Type, cond2Type, then2Type, ...
        // elseType*]

        // Make sure all 'condition' expressions hae type BOOLEAN and all 'then' and
        // an optional 'else' clause have the same type.
        int numCases = argTypes.size() / 2;

        auto &expressionType = argTypes[1];

        for (auto i = 0; i < numCases; i++) {
            auto &conditionType = argTypes[i * 2];
            auto &thenType = argTypes[i * 2 + 1];

            POLLUX_CHECK_EQ(
                conditionType->kind(),
                TypeKind::BOOLEAN,
                "Condition of  SWITCH statement is not bool");

            POLLUX_CHECK(
                *thenType == *expressionType,
                "All then clauses of a SWITCH statement must have the same type. "
                "Expected {}, but got {}.",
                expressionType->toString(),
                thenType->toString());
        }

        bool hasElse = argTypes.size() % 2 == 1;

        if (hasElse) {
            auto &elseClauseType = argTypes.back();

            POLLUX_CHECK(
                *elseClauseType == *expressionType,
                "Else clause of a SWITCH statement must have the same type as 'then' clauses. "
                "Expected {}, but got {}.",
                expressionType->toString(),
                elseClauseType->toString());
        }

        return expressionType;
    }

    TypePtr SwitchCallToSpecialForm::resolveType(
        const std::vector<TypePtr> &argTypes) {
        return SwitchExpr::resolveType(argTypes);
    }

    ExprPtr SwitchCallToSpecialForm::constructSpecialForm(
        const TypePtr &type,
        std::vector<ExprPtr> &&compiledChildren,
        bool /* trackCpuUsage */,
        const core::QueryConfig & /*config*/) {
        bool inputsSupportFlatNoNullsFastPath =
                Expr::allSupportFlatNoNullsFastPath(compiledChildren);
        return std::make_shared<SwitchExpr>(
            type, std::move(compiledChildren), inputsSupportFlatNoNullsFastPath);
    }

    TypePtr IfCallToSpecialForm::resolveType(const std::vector<TypePtr> &argTypes) {
        POLLUX_CHECK(
            argTypes.size() == 3,
            "An IF statement must have 3 clauses, the if clause, the then clause, and the else clause. Expected 3, but got {}.",
            argTypes.size());

        return SwitchCallToSpecialForm::resolveType(argTypes);
    }
} // namespace kumo::pollux::exec
